

import argparse
import time
import os
import torch
from torch import nn
from torch.autograd import Variable
from torch import optim
from torch.utils.data import DataLoader
from dataset import h5Dataset
from pretrain import classifier

#—————————————————————————————————— 1、准备参数 ———————————— ————————————————————— #
use_gpu = torch.cuda.is_available()
epochs = 100
batch_size = 100
n_classes = 10
model_list = ["resnet18","vgg"]
# parse = argparse.ArgumentParser()
# parse.add_argument(
#     '--model',default="vgg", help='vgg, inceptionv3, resnet152')
# parse.add_argument('--phase',default="train", help='train, val')
# parse.add_argument('--n_classes', default=10, type=int)
#
# opt = parse.parse_args()
# print(opt)


# root = '/media/sherlock/Files/kaggle_dog_vs_cat/'
root = "D:\person\learn\pytorch\pytorch_guide\code-of-learn-deep-learning-with-pytorch-master\chapter9_Computer-Vision\kaggle_dog_vs_cat\person_code"

train_list = ['train_feature_{}.hd5f'.format(i) for i in model_list]

val_list = ['val_feature_{}.hd5f'.format(i) for i in model_list]

dataset = {'train': h5Dataset(train_list), 'val': h5Dataset(val_list)}

datasize = {
    'train': dataset['train'].dataset.size(0),
    'val': dataset['val'].dataset.size(0)
}



dataloader = {
    'train':
    DataLoader(
        dataset['train'],
        batch_size=batch_size,
        shuffle=True,
        num_workers=1),
    'val':
    DataLoader(
        dataset['val'],
        batch_size=batch_size,
        shuffle=False,
        num_workers=1)
}

dimension = dataset['train'].dataset.size(1)

mynet = classifier(dimension, n_classes)
mynet.cuda()

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(mynet.parameters(), lr=0.001,weight_decay=1e-5)

def train():
    # train
    for epoch in range(epochs):
        print('{}'.format(epoch + 1))
        print('*' * 10)
        print('Train')
        mynet.train()
        since = time.time()

        running_loss = 0.0
        running_acc = 0.0
        for i, data in enumerate(dataloader['train'], 1):
            feature, label = data
            feature = feature.cuda()
            label = label.cuda()

            # forward
            out = mynet(feature)
            loss = criterion(out, label)
            # backward
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            running_loss += loss.item() * label.size(0)
            _, pred = torch.max(out, 1)
            num_correct = torch.sum(pred == label)
            running_acc += num_correct.item()

            # if i % 50 == 0:
            #     print('Loss: {:.6f}, Acc: {:.6f}'.format(running_loss / (
            #         i * batch_size), running_acc / (i * batch_size)))

        running_loss /= datasize['train']
        running_acc /= datasize['train']
        eplise_time = time.time() - since
        print('Loss: {:.6f}, Acc: {:.6f}, Time: {:.0f}s'.format(
            running_loss, running_acc, eplise_time))
        print('Validation')
        mynet.eval()
        num_correct = 0.0


        eval_loss = 0.0
        for data in dataloader['val']:
            feature, label = data
            feature = Variable(feature, volatile=True).cuda()
            label = Variable(label, volatile=True).cuda()
            # forward
            out = mynet(feature)
            loss = criterion(out, label)

            _, pred = torch.max(out, 1)
            correct = torch.sum(pred == label)
            num_correct += correct.item()
            eval_loss += loss.item() * label.size(0)

        print('Loss: {:.6f}, Acc: {:.6f}'.format(eval_loss / datasize['val'],
                                                 num_correct / datasize['val']))
    print('Finish Training!')

    save_path = os.path.join(root, 'model_save')
    if not os.path.exists(save_path):
        os.mkdir(save_path)

    torch.save(mynet.state_dict(), save_path + '/feature_model.pth')

if __name__ == '__main__':
    train()
